#!/usr/bin/env python
import flowd
import string
import getopt
import sys
import os
import ipcalc

action_types = [
	"src_addr",
	"dst_addr",
	"src_prefix",
	"dst_prefix",
	"src_port",
	"dst_port",
	"exporter",
	"nexthop",
	"src_as",
	"dst_as",
	"protocol",
	"tos",
	"input_interface",
	"output_interface"
]

class tag_class:
	def __init__(self):
		self.name = ""
		self.values = []
		self.match = 1
	def set_name(self, name):
		self.name = name
	def set_match(self, val):
		self.match = val
	def add_value(self, value):
		self.values.append(value)
	def add_values(self, value):
		#self.values.extend(value)
		self.values = value

		
def whatflowfield(field, flow):
	if field == "src_addr" or field == "src_prefix":
		return flow.src_addr
	elif field == "dst_addr" or field == "dst_prefix":
		return flow.dst_addr
	elif field == "src_port":
		return flow.src_port
	elif field == "dst_port":
		return flow.dst_port
	elif field == "src_as":
		return flow.src_as
	elif field == "dst_as":
		return flow.dst_as
	elif field == "protocol":
		return flow.protocol
	elif field == "tos":
		return flow.tos
	elif field == "exporter":
		return flow.agent_addr
	elif field == "nexthop":
		return flow.gateway_addr
	elif field == "input_interface":
		return flow.if_ndx_in
	elif field == "output_interface":
		return flow.if_ndx_out
	return None
	

def usage():
	print >> sys.stderr, "tag.py (flowd.py version %s)" % \
	    flowd.__version__
	print >> sys.stderr, "Usage: tag.py [options] [flowd-store]";
	print >> sys.stderr, "Options:";
	print >> sys.stderr, "      -h       Display this help";
	print >> sys.stderr, "      -t       Path to tagfile";
	print >> sys.stderr, "      -T       Tag number. Hex (in format 0x0) or Decimal";
	#print >> sys.stderr, "Example of tagfile:\n";
	#print >> sys.stderr, "tag-action tag\n type src_addr\n match 192.168.231.1\n match 192.168.231.2\ntag-action tag2\n type src_port\n match 138\ntag-definition 1\n action tag\n action tag2\ntag-definition 2\n action tag2\n";
	#print >> sys.stderr, "  tag-action types: %s"% ', '.join(action_types);
	sys.exit(1);
	
def main():

	try:
		opts, args = getopt.getopt(sys.argv[1:], 't:T:h')
	except getopt.GetoptError:
		print >> sys.stderr, "Invalid commandline arguments"
		usage()

	tagfilepath = None
	tagnum = None
	for o, a in opts:
		if o in ('-h', '--help'):
			usage()
			sys.exit(0)
		if o in ('-t'):
			tagfilepath = a
		if o in ('-T'):
			tagnum = a

	if tagfilepath is None:
		print >> sys.stderr, "No tagfile specified"
		usage()
			
	if tagnum is None:
		print >> sys.stderr, "No tag specified"
		usage()		
			
	if len(args) == 0:
		print >> sys.stderr, "No logfile specified"
		usage()

	#parsing tagfile
	tagfile = open(tagfilepath, "r")
	#tagacthash = {}
	tagactions = []
	line = tagfile.readline()
	while line.strip().startswith("tag-action"):
		action = tag_class()
		tagactname = line.strip().split(" ")[1]
		action.set_name(tagactname)
		line = tagfile.readline()
		#tagactrules = []
		while line.strip().startswith("type"):
			tagacttype = line.strip().split(" ")[1]
			rule = tag_class()
			rule.set_name(tagacttype)
			line = tagfile.readline()
			while line.strip().startswith("match"):
				tagactval = line.strip().split(" ")[1]
				if tagactval == "not":
					rule.set_match(0)
					tagactval = line.strip().split(" ")[2]
				if tagacttype=="src_prefix" or tagacttype=="dst_prefix":
					rule.add_values(ipcalc.Network(tagactval))
				else:
					rule.add_value(tagactval)
				line = tagfile.readline()
			action.add_value(rule)
		tagactions.append(action)
		#tagacthash[tagactname] = tagactrules
	while line.strip().startswith("tag-definition"):
		tagid = line.strip().split(" ")[1]
		tag = tag_class()
		tag.set_name(tagid)
		line = tagfile.readline()
		#tagactnames = []
		while line.strip().startswith("action"):
			for act in tagactions:
				if act.name == line.strip().split(" ")[1]:
					tag.add_value(act)
			line = tagfile.readline()
		if tagid == tagnum:
			break
	
	for act in tag.values:
		for rule in act.values:
			print rule.name, rule.values
	
	if not tag.name.isdigit():
		tag.set_name(long(tag.name,16))
	else:
		tag.set_name(long(tag.name))
	
	#writing tags	
	for ffile in args:
		flogr = flowd.FlowLog(ffile, "rb")
		flogw = flowd.FlowLog("temp", "wb")
		for flow in flogr:
			trueact = 0
			term = 1
			for act in tag.values:
				for rule in act.values:
					if rule.match == 0:
						if str(whatflowfield(rule.name, flow)) not in rule.values:
							trueact = 1
							break
						else: 
							trueact = 0
					else:
						if str(whatflowfield(rule.name, flow)) in rule.values:
							trueact = 1
							break
						else: 
							trueact = 0
				term = term * trueact
			
			if term == 1:		
				if flow.tag == 0:
					flow.fields += 1
				flow.tag = tag.name
			flogw.write_flow(flow)
		
		file = open("temp", "rb")
		buff = file.read()
		file.close()
		file = open(ffile, "wb")
		file.write(buff)
		file.close()
		os.remove("temp")

if __name__ == '__main__': main()